# 意图识别agent

from .base_agent import BaseAgent
from models.state import WorkflowState
from zai import ZhipuAiClient

class IntentRecognitionAgent(BaseAgent):
    def __init__(self):
        self.client = ZhipuAiClient(api_key="db2e03cca9654d4dba9b91b123548374.TftHaYZq4ydQPDjP")
        self.nodes = ["Imperativequery", "ExploratoryAnalysis", "Dataextraction", "other"]
    
    def process(self, state: WorkflowState) -> WorkflowState:
        """
        根据用户输入识别意图类型
        """
        user_input = state.user_input
        
        # 检查是否是澄清询问的回复 - 安全地处理 None 情况
        intent_details = getattr(state, 'intent_details', {}) or {}
        if intent_details.get('waiting_for_clarification', False):
            # 这是用户对澄清询问的回复，直接进行意图识别
            return self._classify_intent(state, user_input, is_clarification_response=True)
        
        # 正常意图识别流程
        return self._classify_intent(state, user_input)
    
    def _classify_intent(self, state: WorkflowState, user_input: str, is_clarification_response: bool = False) -> WorkflowState:
        """核心意图分类逻辑"""
        # 优化后的提示词 - 更加精确和详细
        prompt = """你是一个专业的SQL查询分类专家。请严格按照以下规则对用户问题进行分类：

                    【分类定义和示例】

                    1. Imperativequery（指令式查询）- 用户想要获取一个具体的数值、统计结果或当前状态
                    - 特点：有明确的答案，通常是单个数值或简单列表
                    - 示例：
                        * "显示去年的销售总额" → 返回具体金额
                        * "统计用户总数" → 返回具体数字
                        * "当前库存数量" → 返回具体数量
                        * "最近一周的订单数" → 返回具体数字

                    2. ExploratoryAnalysis（探索式分析）- 用户想要分析数据、发现模式、比较或找出原因
                    - 特点：需要多维度分析，没有单一标准答案
                    - 示例：
                        * "分析销售趋势" → 需要时间趋势分析
                        * "找出最畅销的产品" → 需要排序和比较
                        * "哪个区域业绩最好" → 需要比较分析
                        * "为什么销量下降" → 需要原因分析

                    3. Dataextraction（数据提取）- 用户想要获取原始数据记录、导出完整列表
                    - 特点：获取多条完整记录，用于进一步处理
                    - 示例：
                        * "导出所有客户信息" → 获取完整客户列表
                        * "列出最近三个月的订单" → 获取订单详情列表
                        * "生成客户报表" → 获取结构化数据
                        * "所有产品的详细信息" → 获取完整产品列表

                    4. other（其他）- 与数据查询无关的问题

                    【分类规则】
                    - "找出最畅销的产品" → ExploratoryAnalysis（需要分析比较）
                    - "列出最近三个月的订单" → Dataextraction（获取订单列表）
                    - "显示销售额排名前10的产品" → ExploratoryAnalysis（需要排序分析）
                    - "查询某个客户的订单" → Imperativequery（获取特定信息）

                    请只返回分类名称：Imperativequery、ExploratoryAnalysis、Dataextraction 或 other

                    用户问题："{user_input}"
                    分类结果："""
        
        try:
            # 调用大模型进行意图分类
            response = self.client.chat.completions.create(
                model="glm-4",
                messages=[
                    {"role": "user", "content": prompt.format(user_input=user_input)}
                ],
                temperature=0.1,
                max_tokens=20
            )
            
            # 提取分类结果
            if hasattr(response, 'choices') and len(response.choices) > 0:
                if hasattr(response.choices[0], 'message'):
                    raw_content = response.choices[0].message.content
                    intent_result = self._clean_intent_result(raw_content)
                    
                    # 如果是other意图，生成澄清询问
                    if intent_result == "other" and not is_clarification_response:
                        return self._handle_other_intent(state, user_input)
                    
                    # 更新状态 - 正常意图
                    state.intent = intent_result
                    state.intent_details = {
                        "type": intent_result,
                        "confidence": 0.95,
                        "description": self._get_intent_description(intent_result),
                        "waiting_for_clarification": False,
                        "clarification_question": None
                    }
                else:
                    state = self._set_error_state(state, "响应结构异常")
            else:
                state = self._set_error_state(state, "响应内容为空")
                
        except Exception as e:
            print(f"意图识别出错: {e}")
            state = self._set_error_state(state, str(e))
        
        return state
    
    def _handle_other_intent(self, state: WorkflowState, user_input: str) -> WorkflowState:
        """处理other意图，直接返回固定澄清语句"""
        clarification_question = "对不起，我没能理解你的意图，请告诉我你想查询什么？"
        
        # 设置澄清状态
        return self._set_clarification_state(state, clarification_question)
    
    def _set_clarification_state(self, state: WorkflowState, clarification_text: str) -> WorkflowState:
        """设置澄清状态"""
        state.intent = "other"
        state.intent_details = {
            "type": "other",
            "confidence": 0.1,
            "description": "需要澄清用户意图",
            "waiting_for_clarification": True,
            "clarification_question": clarification_text,
            "clarification_required": True
        }
        return state
    
    def _set_error_state(self, state: WorkflowState, error_msg: str) -> WorkflowState:
        """设置错误状态"""
        state.intent = "other"
        state.intent_details = {
            "type": "other",
            "confidence": 0.1,
            "description": f"分类出错: {error_msg}",
            "error": error_msg,
            "waiting_for_clarification": False,
            "clarification_question": None
        }
        return state
    
    def _clean_intent_result(self, result: str) -> str:
        """清理和验证意图分类结果"""
        if not result:
            return "other"
            
        cleaned = result.strip().lower()
        
        # 直接匹配
        for node in self.nodes:
            if node.lower() in cleaned:
                return node
        
        # 如果不在标准节点中，根据关键词智能匹配
        if any(word in cleaned for word in ['imperative']):
            return "Imperativequery"
        elif any(word in cleaned for word in ['exploratory']):
            return "ExploratoryAnalysis"
        elif any(word in cleaned for word in ['data']):
            return "Dataextraction"
        else:
            return "other"
    
    def _get_intent_description(self, intent_type: str) -> str:
        """获取意图类型的描述信息"""
        descriptions = {
            "Imperativequery": "指令式查询 - 获取具体数据、状态信息或简单统计",
            "ExploratoryAnalysis": "探索式分析 - 数据探索、模式发现、趋势分析",
            "Dataextraction": "数据提取 - 记录查询、数据导出、报表生成",
            "other": "其他问题 - 不属于SQL查询范畴"
        }
        return descriptions.get(intent_type, "未知意图类型")